package com.equipment.training.web.interceptor;

import com.alibaba.fastjson.JSON;
import com.equipment.training.web.annotation.ApiEncrypt;
import com.equipment.training.web.security.RSAUtils;
import com.equipment.training.web.security.RsaHelper;
import org.springframework.core.MethodParameter;
import org.springframework.core.annotation.Order;
import org.springframework.http.MediaType;
import org.springframework.http.server.ServerHttpRequest;
import org.springframework.http.server.ServerHttpResponse;
import org.springframework.util.ResourceUtils;
import org.springframework.web.bind.annotation.ControllerAdvice;
import org.springframework.web.servlet.mvc.method.annotation.ResponseBodyAdvice;

import java.io.File;
import java.io.FileInputStream;
import java.io.InputStream;
import java.security.PublicKey;

/**
 * ${DESCRIPTION}
 *
 * @author 杨凯
 * @create 2018-03-27 12:44
 **/
@Order(1)
@ControllerAdvice
public class MyResponseBodyAdvice implements ResponseBodyAdvice {

	//是否加密
	private boolean encode = false;

	@Override
	public boolean supports(MethodParameter methodParameter, Class aClass) {
		//这里可以根据自己的需求
		return true;
	}

	@Override
	public Object beforeBodyWrite(Object o, MethodParameter methodParameter, MediaType mediaType, Class aClass, ServerHttpRequest serverHttpRequest, ServerHttpResponse serverHttpResponse) {

		//判断返回的对象是单个对象，还是list，活着是map
		if (o == null) {
			return null;
		}
		if (methodParameter.getMethod().isAnnotationPresent(ApiEncrypt.class)) {
			ApiEncrypt apiEncrypt = methodParameter.getMethodAnnotation(ApiEncrypt.class);
			//是否加密
			encode = apiEncrypt.encode();
			if (!encode) {
				return o;
			}
			Object retObj = null;
			//Single Object
			retObj = handleSingleObject(o);
			return retObj;
		}

		return o;
	}

	/**
	 * 处理返回值是单个enity对象
	 *
	 * @param o
	 * @return
	 */
	private Object handleSingleObject(Object o) {
		return getNewVal(o);
	}

	/**
	 * 获取加密后的新值
	 *
	 * @param obj
	 * @return
	 */
	private String getNewVal(Object obj) {
		String newVal = "";
		try {
			if (obj != null) {
				String json = JSON.toJSONString(obj);
				System.out.println(json);
				if (encode) {
					File file = ResourceUtils.getFile("classpath:rsa_public_key.pem");
					// 从文件中得到公钥
					InputStream inPublic = new FileInputStream(file);
					PublicKey publicKey = RSAUtils.loadPublicKey(inPublic);
					// 加密
					//加密后Key
					newVal = RsaHelper.encryptDataFromStr(json, publicKey);
					System.out.println(newVal);
				} else {
					newVal = json;
				}
			}
		} catch (Exception e) {
			e.printStackTrace();
		}

		return newVal;
	}
}
